#include <wdm.h>

#include "CpuStructs.h"

#pragma warning(disable : 4100)

enum SegmentDescriptorTableType
{
	SDT_TYPE_GDT,
	SDT_TYPE_LDT
};

VOID MEMAPPER_Unload(_In_ struct _DRIVER_OBJECT *DriverObject);
NTSTATUS MEMAPPER_DispatchCreate(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							  _Inout_ struct _IRP *Irp);
NTSTATUS MEMAPPER_DispatchClose(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							 _Inout_ struct _IRP *Irp);
NTSTATUS MEMAPPER_DispatchRead(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							_Inout_ struct _IRP *Irp);
NTSTATUS MEMAPPER_DispatchWrite(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							 _Inout_ struct _IRP *Irp);
VOID MEMAPPER_GetSegmentDescriptorTable(enum SegmentDescriptorTableType type,
										segment_descriptor_t** table,
										PUINT16 table_limit);

#define DRIVER_NAME "MEMAPPER"
#define NT_DEVICE_NAME L"\\Device\\MEMAPPER"
#define DOS_DEVICE_NAME L"\\DosDevices\\MEMAPPER"

#define BUFFER_SIZE 1024
char buffer[BUFFER_SIZE];


/*
http://msdn.microsoft.com/en-us/library/windows/hardware/ff544113%28v=vs.85%29.aspx
*/
DRIVER_INITIALIZE DriverEntry;
NTSTATUS DriverEntry(_In_ struct _DRIVER_OBJECT *DriverObject,
					 _In_ PUNICODE_STRING RegistryPath)
{
	/* Using NTSTATUS values:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff565436%28v=vs.85%29.aspx
	*/
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING DeviceName;
	UNICODE_STRING DosDeviceName;
	PDEVICE_OBJECT DeviceObject = NULL;

	DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
				DRIVER_NAME": --> DriverEntry\n");
	
	/* Setup driver object */
	DriverObject->MajorFunction[IRP_MJ_CREATE] = MEMAPPER_DispatchCreate;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = MEMAPPER_DispatchClose;
	DriverObject->MajorFunction[IRP_MJ_READ] = MEMAPPER_DispatchRead;
	DriverObject->MajorFunction[IRP_MJ_WRITE] = MEMAPPER_DispatchWrite;

	DriverObject->DriverUnload = MEMAPPER_Unload;

	/* Create the device */
	RtlInitUnicodeString(&DeviceName, NT_DEVICE_NAME);
	RtlInitUnicodeString(&DosDeviceName, DOS_DEVICE_NAME);
	/*
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff548397%28v=vs.85%29.aspx
	*/
	status = IoCreateDevice(DriverObject, 0, &DeviceName,
							FILE_DEVICE_UNKNOWN,
							FILE_DEVICE_SECURE_OPEN,
							FALSE,
							&DeviceObject);
	if (!NT_SUCCESS(status))
	{
		/*
		http://msdn.microsoft.com/en-us/library/windows/hardware/ff543632%28v=vs.85%29.aspx
		*/
		//DbgPrint("Failed to IoCreateDevice\n");
		/*
		http://msdn.microsoft.com/en-us/library/windows/hardware/ff543634%28v=vs.85%29.aspx
		*/
		DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL,
					DRIVER_NAME": Error creating device (Failed to IoCreateDevice)\n");
		return status;
	}

	// Map the "DOS Device Name" to the "NT Device Name"
	status = IoCreateSymbolicLink(&DosDeviceName, &DeviceName);
	if (!NT_SUCCESS(status))
	{
		DbgPrintEx(DPFLTR_IHVDRIVER_ID,	DPFLTR_ERROR_LEVEL,
					DRIVER_NAME": Error creating symlink (Failed to IoCreateSymbolicLink)\n");
		return status;
	}

	// Initializing a Device Object
	// http://msdn.microsoft.com/en-us/library/windows/hardware/ff547807(v=vs.85).aspx
	DeviceObject->Flags |= DO_BUFFERED_IO;
	// must be cleared if DeviceObject was created outside of the DriverEntry function
	// DeviceObject->Flags &= (~DO_DEVICE_INITIALIZING);

	KdPrintEx((DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
				DRIVER_NAME": <-- DriverEntry\n"));
	
	return status;
}

/*
http://msdn.microsoft.com/en-us/library/windows/hardware/ff564886%28v=vs.85%29.aspx
*/
DRIVER_UNLOAD MEMAPPER_Unload;
VOID MEMAPPER_Unload(_In_ struct _DRIVER_OBJECT *DriverObject)
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING DosDriverName;

	KdPrintEx((DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
				DRIVER_NAME": --> Unload\n"));

	RtlInitUnicodeString(&DosDriverName, DOS_DEVICE_NAME);
	status = IoDeleteSymbolicLink(&DosDriverName);

	IoDeleteDevice(DriverObject->DeviceObject);

	DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
				DRIVER_NAME": <-- Unload\n");
}

DRIVER_DISPATCH MEMAPPER_DispatchCreate;
NTSTATUS MEMAPPER_DispatchCreate(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							  _Inout_ struct _IRP *Irp)
{
	/* Controlling Device Namespace Access:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff542068(v=vs.85).aspx
	fail all IRP_MJ_CREATE requests that specify
	an IrpSp->FileObject->FileName parameter whose length is nonzero,
	because we don't support namespaces
	*/
	return STATUS_SUCCESS;
}

DRIVER_DISPATCH MEMAPPER_DispatchClose;
NTSTATUS MEMAPPER_DispatchClose(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							 _Inout_ struct _IRP *Irp)
{
	return STATUS_SUCCESS;
}

/*
http://msdn.microsoft.com/en-us/library/windows/hardware/ff543376%28v=vs.85%29.aspx
*/
DRIVER_DISPATCH MEMAPPER_DispatchRead;
NTSTATUS MEMAPPER_DispatchRead(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							_Inout_ struct _IRP *Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	PIO_STACK_LOCATION stack = NULL;
	LARGE_INTEGER offset;
	ULONG length;
	PVOID client_buffer;
	ULONG bytes_to_read;

	stack = IoGetCurrentIrpStackLocation(Irp);

	offset = stack->Parameters.Read.ByteOffset;
	length = stack->Parameters.Read.Length;
	client_buffer = Irp->AssociatedIrp.SystemBuffer;

	/* Perform the request */
	if (offset.QuadPart < BUFFER_SIZE)
	{
		/* read as much bytes as we can */
		bytes_to_read = (offset.QuadPart + length <= BUFFER_SIZE)
							? length
							: BUFFER_SIZE - (ULONG)offset.QuadPart;
		RtlCopyMemory(client_buffer, buffer + offset.QuadPart, bytes_to_read);
	}
	else
	{
		status = STATUS_UNSUCCESSFUL;
		bytes_to_read = 0;
	}

	/* Complete the request*/
	/* IO_STATUS_BLOCK structure:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff550671%28v=vs.85%29.aspx
	*/
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = bytes_to_read;

	/* IoCompleteRequest routine:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff548343%28v=vs.85%29.aspx
	*/
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return STATUS_SUCCESS;
}

/* Assuming BUFFERED_IO */
DRIVER_DISPATCH MEMAPPER_DispatchWrite;
NTSTATUS MEMAPPER_DispatchWrite(_Inout_ struct _DEVICE_OBJECT *DeviceObject,
							 _Inout_ struct _IRP *Irp)
{
	/*
	IRP:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff550694%28v=vs.85%29.aspx
	IO_STACK_LOCATION:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff550659%28v=vs.85%29.aspx
	*/
	NTSTATUS status = STATUS_SUCCESS;
	PIO_STACK_LOCATION stack = NULL;
	
	segment_descriptor_t* table = NULL;
	UINT16 table_limit = 0;
	
	LARGE_INTEGER offset;
	ULONG length;
	PVOID client_buffer;
	ULONG bytes_to_write;
	
	/* Get request parameters */
	/* IoGetCurrentIrpStackLocation routine:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff549174%28v=vs.85%29.aspx
	*/
	stack = IoGetCurrentIrpStackLocation(Irp);

	offset = stack->Parameters.Write.ByteOffset;
	length = stack->Parameters.Write.Length;
	client_buffer = Irp->AssociatedIrp.SystemBuffer;

	if (length > 0)
	{
		enum SegmentDescriptorTableType table_type = ((enum SegmentDescriptorTableType*) client_buffer)[0];
		MEMAPPER_GetSegmentDescriptorTable(table_type, &table, &table_limit);

		bytes_to_write = min((table_limit * sizeof(segment_descriptor_t)), (BUFFER_SIZE));
		RtlCopyMemory(buffer, table, bytes_to_write);
	}
	else
	{
		status = STATUS_UNSUCCESSFUL;
		bytes_to_write = 0;
	}

	/* Complete the request*/
	/* IO_STATUS_BLOCK structure:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff550671%28v=vs.85%29.aspx
	*/
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = bytes_to_write;

	/* IoCompleteRequest routine:
	http://msdn.microsoft.com/en-us/library/windows/hardware/ff548343%28v=vs.85%29.aspx
	*/
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return status;
}

VOID MEMAPPER_PrintSegmentDescriptorTable(segment_descriptor_t* table,
										  UINT16 table_limit)
{
	UINT16 i;
	for (i = 0; i < table_limit; ++i)
	{
		segment_descriptor_t entry = table[i];
		DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
			DRIVER_NAME": "
			"Higher Half: (base_31:24 = %x, granularity = %x, \n"
			"default_operation_size = %x, avl = %x, \n"
			"segment_limit_19:16 = %x, segment_present = %x, \n"
			"privilege_level = %x, descriptor_type = %x, \n"
			"segment_type = %x, base_address_23:16 = %hx)\n"
			"Lower Half: (segment_limit_15:00 = %hx, base_address_15:00 = %hx)\n",
			entry.base_2431, entry.granularity,
			entry.default_operation_size, entry.avl,
			entry.segment_limit_1619, entry.segment_present,
			entry.privilege_level, entry.descriptor_type,
			entry.segment_type, entry.base_address_1623,
			entry.segment_limit_0015, entry.base_address_0015);
	}
}

VOID MEMAPPER_GDT(segment_descriptor_t** table,
				  PUINT16 table_limit)
{
	system_table_register_t gdtr = {0xcafe, 0xdeadcafe};
	_asm sgdt gdtr;
	DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
		DRIVER_NAME": gdtr.base = %I32x, gdtr.limit = %I32x\n",
		gdtr.base, gdtr.limit);

	*table = (segment_descriptor_t*) gdtr.base;
	*table_limit = gdtr.limit;
}

VOID MEMAPPER_LDT(segment_descriptor_t** table,
				  PUINT16 table_limit)
{
	union segment_selector_t ldtr_seg_sel;
	segment_descriptor_t* ldtr_segment_descr = NULL;
	_asm sldt ldtr_seg_sel;
	DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_TRACE_LEVEL,
		DRIVER_NAME": ldtr segment selector = %hx "
		"(rpl = %x, table = %x, index = %hx)\n",
		ldtr_seg_sel.selector, ldtr_seg_sel.rpl,
		ldtr_seg_sel.table, ldtr_seg_sel.index);

	MEMAPPER_GDT(table, table_limit);
	ldtr_segment_descr = &(*table)[ldtr_seg_sel.index];
}

VOID MEMAPPER_GetSegmentDescriptorTable(enum SegmentDescriptorTableType type,
										segment_descriptor_t** table,
										PUINT16 table_limit)
{
	switch (type)
	{
	case SDT_TYPE_LDT:
		MEMAPPER_LDT(table, table_limit);
		break;
	default:
		MEMAPPER_GDT(table, table_limit);
	};
	MEMAPPER_PrintSegmentDescriptorTable(*table, *table_limit);
}
